package com.simafei.flow.core.data.ds;

import lombok.NonNull;
import org.springframework.jdbc.datasource.AbstractDataSource;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 动态数据源
 *
 * @author mcex
 */
public class FlowDataSource extends AbstractDataSource {

    private final Map<String, DataSource> dataSourceMap;

    public FlowDataSource() {
        this.dataSourceMap = new ConcurrentHashMap<>();
    }

    public void add(String key, DataSource dataSource) {
        dataSourceMap.put(key, dataSource);
    }

    public void remove(String key) {
        dataSourceMap.remove(key);
    }

    @Override
    public Connection getConnection() throws SQLException {
        return determineDataSource().getConnection();
    }

    @Override
    public Connection getConnection(String username, String password) throws SQLException {
        return determineDataSource().getConnection(username, password);
    }

    public DataSource determineDataSource() {
        String dsKey = FlowDataSourceContextHolder.get();
        return getDataSource(dsKey);
    }

    public DataSource getDataSource(@NonNull String ds) {
        DataSource dataSource = dataSourceMap.get(ds);
        if (Objects.nonNull(dataSource)) {
            return dataSource;
        } else {
            throw new IllegalArgumentException("dataSource not found: " + ds);
        }
    }

    public List<String> getDataSourceNames() {
        return dataSourceMap.keySet().stream().sorted().toList();
    }

    @Override
    @SuppressWarnings("unchecked")
    public <T> T unwrap(Class<T> face) throws SQLException {
        if (face.isInstance(this)) {
            return (T) this;
        }
        return determineDataSource().unwrap(face);
    }

    @Override
    public boolean isWrapperFor(Class<?> face) throws SQLException {
        return (face.isInstance(this) || determineDataSource().isWrapperFor(face));
    }

    public boolean exists(String key) {
        return dataSourceMap.containsKey(key);
    }
}